import librosa
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import os
# 设置环境变量 KMP_DUPLICATE_LIB_OK=TRUE 作为临时解决方案
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

def extract_mfcc(file_path):
    y, sr = librosa.load(file_path, sr=None)
    mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
    mfcc = np.mean(mfcc.T, axis=0)
    return mfcc

def extract_mfcc_from_folder(folder_path):
    all_mfcc = []
    labels = []
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.lower().endswith(('.wav', '.mp3')):  # 确保是音频文件
                file_path = os.path.join(root, file)
                mfcc = extract_mfcc(file_path)
                all_mfcc.append(mfcc)
                # 假设文件夹名称为“你”或“好”，作为标签
                label = 0 if '你' in root else 1
                labels.append(label)
    return all_mfcc, labels

class SpeechDataset(Dataset):
    def __init__(self, mfcc_features, labels):
        self.mfcc_features = mfcc_features
        self.labels = labels

    def __len__(self):
        return len(self.mfcc_features)

    def __getitem__(self, idx):
        feature = torch.tensor(self.mfcc_features[idx], dtype=torch.float32)
        label = torch.tensor(self.labels[idx], dtype=torch.long)
        return feature, label

# 训练集文件夹路径，里面分别有'你'和'好'的子文件夹，子文件夹中存放对应音频
train_folder_path = r'E:\实训\音频数据集'
# 预测的音频所在文件夹路径
predict_folder_path = r'E:\实训\测试数据集'

# 提取训练集的MFCC特征和标签
mfcc_features_train, labels_train = extract_mfcc_from_folder(train_folder_path)

dataset_train = SpeechDataset(mfcc_features_train, labels_train)
dataloader_train = DataLoader(dataset_train, batch_size=2, shuffle=True)

# 构建模型
class CNNModel(nn.Module):
    def __init__(self):
        super(CNNModel, self).__init__()
        self.conv1 = nn.Conv1d(1, 32, kernel_size=3, stride=1)
        self.pool = nn.MaxPool1d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(32 * 5, 64)
        self.fc2 = nn.Linear(64, 2)

    def forward(self, x):
        x = x.unsqueeze(1)  # 添加 channel 维度
        x = self.pool(F.relu(self.conv1(x)))
        x = x.view(-1, 32 * 5)  # 展平
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = CNNModel()

# 训练模型
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10
for epoch in range(num_epochs):
    for inputs, labels in dataloader_train:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {loss.item()}')

# 提取预测文件夹中音频的MFCC特征
mfcc_features_predict, _ = extract_mfcc_from_folder(predict_folder_path)

# 进行预测
for mfcc in mfcc_features_predict:
    inputs = torch.tensor(mfcc, dtype=torch.float32).unsqueeze(0)  # 添加 batch 维度
    outputs = model(inputs)
    _, predicted = torch.max(outputs, 1)
    print(f'Prediction: {"你" if predicted.item() == 0 else "好"}')